import  os
import json
from lxml import etree as ET
from xml.dom import minidom

def edit_xml(categories, objects, imgName, dir):
    save_xml_path = os.path.join(dir, "%s.xml" % imgName)  # xml
    #save_xml_path = "VOCdevkit/VOC2007/Annotations/2.xml"

    root = ET.Element("annotation")
    # root.set("version", "1.0")  
    folder = ET.SubElement(root, "folder")
    folder.text = "none"
    filename = ET.SubElement(root, "filename")
    filename.text = "none"
    source = ET.SubElement(root, "source")
    source.text = "201908"
    owner = ET.SubElement(root, "owner")
    owner.text = "YZN"
    size = ET.SubElement(root, "size")
    width = ET.SubElement(size, "width")
    width.text = str(640)
    height = ET.SubElement(size, "height")
    height.text = str(512)
    depth = ET.SubElement(size, "depth")
    depth.text = "3"
    segmented = ET.SubElement(root, "segmented")
    segmented.text = "0"
    for obj in objects:  #  
        object = ET.SubElement(root, "object")
        name = ET.SubElement(object, "name")  # number
        name.text = categories[obj["category_id"] - 1]["name"]
        # meaning = ET.SubElement(object, "meaning")  # name
        # meaning.text = inf_value[0]
        pose = ET.SubElement(object, "pose")
        pose.text = "Unspecified"
        truncated = ET.SubElement(object, "truncated")
        truncated.text = "0"
        difficult = ET.SubElement(object, "difficult")
        difficult.text = "0"
        bndbox = ET.SubElement(object, "bndbox")
        xmin = ET.SubElement(bndbox, "xmin")
        xmin.text = str(int(obj["bbox"][0]))
        ymin = ET.SubElement(bndbox, "ymin")
        ymin.text = str(int(obj["bbox"][1]))
        xmax = ET.SubElement(bndbox, "xmax")
        xmax.text = str(int(obj["bbox"][0]) + int(obj["bbox"][2]))
        ymax = ET.SubElement(bndbox, "ymax")
        ymax.text = str(int(obj["bbox"][1]) + int(obj["bbox"][3]))
    #print(ET.tostring(root, encoding='UTF-8'))
    tree = ET.ElementTree(root)
    tree.write(save_xml_path, encoding="UTF-8", xml_declaration=True)
    root = ET.parse(save_xml_path) 
    file_lines = minidom.parseString(ET.tostring(root, encoding="Utf-8")).toprettyxml(
        indent="\t") 
    file_line = open(save_xml_path, "w", encoding="utf-8")  
    file_line.write(file_lines)
    file_line.close() 

filedir = "VOCdevkit/VOC2007/Json/thermal_annotations.json"
annos = json.loads(open(filedir).read())

categories = annos["categories"]
annotations = annos["annotations"]
images = annos["images"]

beginId = 0
for imgId in range(0, len(images)):
	#  json 中的ID图片有待检测目标，且该id图片在 train文件夹中
    objects = []
    for annId in range(beginId, len(annotations)):
        if imgId == annotations[annId]["image_id"]:
            objects.append(annotations[annId])
        else :
            beginId = annId
            break
    imgName = images[imgId]["file_name"][14:24]
    edit_xml(categories,objects, imgName, dir = "VOCdevkit/VOC2007/Annotations")
print("successful!")
        

